#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#中国地震情况统计
#爬取中国地震台网数据
import os.path
import sys
import json

import requests
from pyecharts.charts import Geo
from pyecharts import options as opts
from pyecharts.globals import GeoType, ThemeType

def getData():
    r = requests.get('http://www.ceic.ac.cn/ajax/search?page=1&&start=&&end=&&jingdu1=&&jingdu2=&&weidu1=&&weidu2=&&height1=&&height2=&&zhenji1=&&zhenji2=')
    if r.status_code == requests.codes.ok:
        data = json.loads(r.text[1:-1])['shuju']
        for x in range(json.loads(r.text[1:-1])['num']-1):
            r = requests.get('http://www.ceic.ac.cn/ajax/search?page=' + str(x+2) + '&&start=&&end=&&jingdu1=&&jingdu2=&&weidu1=&&weidu2=&&height1=&&height2=&&zhenji1=&&zhenji2=')
            data += json.loads(r.text[1:-1])['shuju']
    with open('earthquake_data.json', 'w') as f:
        json.dump(data, f, sort_keys=True, indent=4)

def drawHeatmap(chinaOnly=False, round_=None, deep=3):
    #生成地理热力图
    #chinaOnly: bool = True True为仅显示中国地区，否则显示全世界
    #round_: int = None 当存在时，对经纬度进行四舍五入并合并，降低精度以减少生成文件大小，值为保留的小数位数，可选0、1、2
    #deep: int/float = 颜色深度，数值越大，颜色越深
    chart = Geo(init_opts=opts.InitOpts(theme=ThemeType.CHALK))
    if chinaOnly:
        chart.add_schema(maptype='china')
    else:
        chart.add_schema(maptype='world')
    dataPair = []
    if round_:
        heatmapData = {}
        for x in data:
            longitude = round(float(x['EPI_LON']), round_)
            latitude = round(float(x['EPI_LAT']), round_)
            heatmapData[str(longitude)+'-'+str(latitude)] = [longitude, latitude, heatmapData.get(str(longitude)+str(latitude), [None, None, 0])[2]+1]
        for k,v in heatmapData.items():
            dataPair.append((k, v[2]*deep))
            chart.add_coordinate(k, v[0], v[1])
    else:
        print(1)
        for x in data:
            dataPair.append((x['id'], deep))
            chart.add_coordinate(x['id'], x['EPI_LON'], x['EPI_LAT'])
    chart.add('地震频率', dataPair, type_=GeoType.HEATMAP)
    chart.set_global_opts(visualmap_opts=opts.VisualMapOpts())
    chart.render()

if __name__ == '__main__':
    if os.path.isfile('earthquake_data.json'):
        with open('earthquake_data.json', 'r') as f:
            data = json.load(f)
    else:
        getData()
    drawHeatmap()